import { WebSocket } from 'ws'
import { Client } from 'ssh2'

const os = require('os')
const pty = require('node-pty')

const wsMap = new Map()
const sessionMap = new Map()
const localMap = new Map()

export const initWS = () => {
  const client = new WebSocket.Server({ port: 10095 })
  client.on('connection', (ws) => {
    ws.on('message', (msg) => {
      const data = JSON.parse(msg)
      if (data.type === 'ssh-conn') {
        wsMap.set(data.id, ws)
        connSSH(ws, data)
      }
      if (data.type === 'local-conn') {
        wsMap.set(data.id, ws)
        connLoacl(ws, data)
      }
      if (data.type === 'ssh-cmd') {
        handleCmd(data)
      }
      if (data.type === 'local-cmd') {
        handleLocalCmd(data)
      }
      if (data.type === 'local-resize') {
        handleLocalResize(data)
      }
      if (data.type === 'ssh-resize') {
        handleSSHResize(data)
      }
    })
    ws.on('close', () => {
      console.log('Client disconnected')
      wsMap.forEach((value, key, map) => {
        if (value === ws) {
          map.delete(key)
        }
      })
      localMap.forEach((value, key, map) => {
        if (value === ws) {
          value.kill()
          map.delete(key)
        }
      })
    })
  })
}

const connSSH = (ws, data) => {
  const sshClient = initSSH(ws, data)
  sshClient.on('ready', () => {
    console.log('Connected to remote server.')
    const session = sshClient.shell(
      { term: 'xterm-256color', cols: data.cols, rows: data.rows },
      (err, stream) => {
        if (err) {
          throw err
        }
        sessionMap.set(data.id, stream)

        stream.on('close', (code, signal) => {
          console.log(`ssh session close exitCode: ${code}, signal: ${signal}`)
          sshClient.end()
          sessionMap.forEach((value, key, map) => {
            if (value === session) {
              map.delete(key)
            }
          })
        })
        stream.on('data', (data) => {
          const msgData = {
            type: 'ssh-msg',
            data: String(data)
          }
          ws.send(JSON.stringify(msgData))
        })
      }
    )
  })
}

const initSSH = (ws, data) => {
  const client = new Client()
  client.on('error', (err) => {
    console.log(`ssh client error:${err.message}`)
    const sshError = {
      type: 'ssh-error',
      data: err.message
    }
    ws.send(JSON.stringify(sshError))
  })
  client.connect({
    host: data.host,
    port: data.port,
    username: data.username,
    password: data.password
  })
  return client
}

const handleCmd = (data) => {
  const session = sessionMap.get(data.id)
  session.write(data.cmd)
}

const connLoacl = (ws, data) => {
  const shell = getShell()
  const ptyProcess = pty.spawn(shell, [], {
    name: 'xterm-color',
    cols: data.cols,
    rows: data.rows,
    cwd: process.env.HOME,
    env: process.env
  })
  localMap.set(data.id, ptyProcess)
  ptyProcess.onData((data) => {
    const msgData = {
      type: 'local-msg',
      data: data
    }
    ws.send(JSON.stringify(msgData))
  })
}

const handleLocalCmd = (data) => {
  const ptyProcess = localMap.get(data.id)
  ptyProcess.write(data.cmd)
}

const handleLocalResize = (data) => {
  const ptyProcess = localMap.get(data.id)
  ptyProcess.resize(data.cols, data.rows)
}

const handleSSHResize = (data) => {
  const session = sessionMap.get(data.id)
  session.write(`stty cols ${data.cols} rows ${data.cols} \r`)
}

const getShell = () => {
  if (os.platform() === 'win32') {
    return 'powershell.exe'
  }

  if (os.platform() === 'darwin') {
    return '/bin/zsh'
  }

  return '/bin/bash'
}
